{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先定义工具函数，用来完成Agent的任务。注意：大模型会根据函数的注释来判断使用哪个函数来完成任务。所以，注释一定要写清楚函数的功能和返回值。\n",
    "\n",
    "然后把工具函数放入FunctionTool对象中，供Agent能够使用。\n",
    "\n",
    "用 LlamaIndex 实现一个简单的 agent demo 比较容易，LlamaIndex 实现 Agent 需要导入 ReActAgent 和 Function Tool。\n",
    "\n",
    "ReActAgent 是什么？\n",
    "\n",
    "ReActAgent 通过结合推理（Reasoning）和行动（Acting）来创建动态的 LLM Agent 的框架。该方法允许 LLM 模型通过在复杂环境中交替进行推理步骤和行动步骤来更有效地执行任务。ReActAgent 将推理和动作形成了闭环，Agent 可以自己完成给定的任务。\n",
    "\n",
    "一个典型的 ReActAgent 遵循以下循环：\n",
    "\n",
    "初始推理：代理首先进行推理步骤，以理解任务、收集相关信息并决定下一步行为。\n",
    "行动：代理基于其推理采取行动——例如查询API、检索数据或执行命令。\n",
    "观察：代理观察行动的结果并收集任何新的信息。\n",
    "优化推理：利用新信息，代理再次进行推理，更新其理解、计划或假设。\n",
    "重复：代理重复该循环，在推理和行动之间交替，直到达到满意的结论或完成任务。\n",
    "\n",
    "实现最简单的代码，通过外部工具做算术题，只是一个简单的例子。这个如果不用 Agent，其实大模型也可以回答。看一下具体的代码实现：\n",
    "\n",
    "\n",
    "首先准备各种key和模型名称"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# 加载环境变量\n",
    "load_dotenv()\n",
    "# 从环境变量中读取api_key\n",
    "api_key = os.getenv('ZISHU_API_KEY')\n",
    "base_url = \"http://101.132.164.17:8000/v1\"\n",
    "chat_model = \"glm-4-flash\"\n",
    "emb_model = \"embedding-3\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后来构建llm，其实任何能用的llm都行。这里自定义一个。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "from pydantic import Field  # 导入Field，用于Pydantic模型中定义字段的元数据\n",
    "from llama_index.core.llms import (\n",
    "    CustomLLM,\n",
    "    CompletionResponse,\n",
    "    LLMMetadata,\n",
    ")\n",
    "from llama_index.core.embeddings import BaseEmbedding\n",
    "from llama_index.core.llms.callbacks import llm_completion_callback\n",
    "from typing import List, Any, Generator\n",
    "# 定义OurLLM类，继承自CustomLLM基类\n",
    "class OurLLM(CustomLLM):\n",
    "    api_key: str = Field(default=api_key)\n",
    "    base_url: str = Field(default=base_url)\n",
    "    model_name: str = Field(default=chat_model)\n",
    "    client: OpenAI = Field(default=None, exclude=True)  # 显式声明 client 字段\n",
    "\n",
    "    def __init__(self, api_key: str, base_url: str, model_name: str = chat_model, **data: Any):\n",
    "        super().__init__(**data)\n",
    "        self.api_key = api_key\n",
    "        self.base_url = base_url\n",
    "        self.model_name = model_name\n",
    "        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)  # 使用传入的api_key和base_url初始化 client 实例\n",
    "\n",
    "    @property\n",
    "    def metadata(self) -> LLMMetadata:\n",
    "        \"\"\"Get LLM metadata.\"\"\"\n",
    "        return LLMMetadata(\n",
    "            model_name=self.model_name,\n",
    "        )\n",
    "\n",
    "    @llm_completion_callback()\n",
    "    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:\n",
    "        response = self.client.chat.completions.create(model=self.model_name, messages=[{\"role\": \"user\", \"content\": prompt}])\n",
    "        if hasattr(response, 'choices') and len(response.choices) > 0:\n",
    "            response_text = response.choices[0].message.content\n",
    "            return CompletionResponse(text=response_text)\n",
    "        else:\n",
    "            raise Exception(f\"Unexpected response format: {response}\")\n",
    "\n",
    "    @llm_completion_callback()\n",
    "    def stream_complete(\n",
    "        self, prompt: str, **kwargs: Any\n",
    "    ) -> Generator[CompletionResponse, None, None]:\n",
    "        response = self.client.chat.completions.create(\n",
    "            model=self.model_name,\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            stream=True\n",
    "        )\n",
    "\n",
    "        try:\n",
    "            for chunk in response:\n",
    "                chunk_message = chunk.choices[0].delta\n",
    "                if not chunk_message.content:\n",
    "                    continue\n",
    "                content = chunk_message.content\n",
    "                yield CompletionResponse(text=content, delta=content)\n",
    "\n",
    "        except Exception as e:\n",
    "            raise Exception(f\"Unexpected response format: {e}\")\n",
    "\n",
    "llm = OurLLM(api_key=api_key, base_url=base_url, model_name=chat_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测试一下这个llm能用吗？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "我是一个名为 ChatGLM 的人工智能助手，是基于清华大学 KEG 实验室和智谱 AI 公司于 2024 年共同训练的语言模型开发的。我的任务是针对用户的问题和要求提供适当的答复和支持。"
     ]
    }
   ],
   "source": [
    "response = llm.stream_complete(\"你是谁？\")\n",
    "for chunk in response:\n",
    "    print(chunk, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 72d00574-e60d-4b7e-b4eb-5347a165a443. Step input: 20+（2*4）等于多少？使用工具计算每一步\n",
      "\u001b[1;3;38;5;200mThought: The user is asking for a calculation involving addition and multiplication. I will use the provided tools to calculate the result step by step.\n",
      "Action: multiply\n",
      "Action Input: {'a': 2, 'b': 4}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 8\n",
      "\u001b[0m> Running step a56ba38f-ff4f-4a45-b80e-d67f4aae6304. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: (user's language). I need to use a tool to help me answer the question.\n",
      "Action: add\n",
      "Action Input: {'a': 20, 'b': 8}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 28\n",
      "\u001b[0m> Running step 10687254-b221-4c1f-aa3d-6f6def97271d. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: 20 + (2 * 4) = 28\n",
      "\u001b[0m20 + (2 * 4) = 28\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "from llama_index.core.agent import ReActAgent\n",
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def multiply(a: float, b: float) -> float:\n",
    "    \"\"\"Multiply two numbers and returns the product\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: float, b: float) -> float:\n",
    "    \"\"\"Add two numbers and returns the sum\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def main():\n",
    "\n",
    "    multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "    add_tool = FunctionTool.from_defaults(fn=add)\n",
    "\n",
    "    # 创建ReActAgent实例\n",
    "    agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)\n",
    "\n",
    "    response = agent.chat(\"20+（2*4）等于多少？使用工具计算每一步\")\n",
    "\n",
    "    print(response)\n",
    "\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当我们问大模型一个天气的问题，当没有工具时，大模型这么回答，作为大语言模型，他不知道天气情况并给出去哪里可以查到天气情况。\n",
    "现在为我们的 Agent 添加一个查询天气的方法，返回假数据做测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step d5b3667d-895d-4c64-b21b-0e9748d61d73. Step input: 纽约天气怎么样?\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: Chinese. I need to use a tool to help me answer the question.\n",
      "Action: get_weather\n",
      "Action Input: {'city': 'NY'}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 20\n",
      "\u001b[0m> Running step 83986747-926a-4bd3-b19a-b49059163175. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: 纽约的天气是20度。\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "def get_weather(city: str) -> int:\n",
    "    \"\"\"\n",
    "    Gets the weather temperature of a specified city.\n",
    "\n",
    "    Args:\n",
    "    city (str): The name or abbreviation of the city.\n",
    "\n",
    "    Returns:\n",
    "    int: The temperature of the city. Returns 20 for 'NY' (New York),\n",
    "         30 for 'BJ' (Beijing), and -1 for unknown cities.\n",
    "    \"\"\"\n",
    "\n",
    "    # Convert the input city to uppercase to handle case-insensitive comparisons\n",
    "    city = city.upper()\n",
    "\n",
    "    # Check if the city is New York ('NY')\n",
    "    if city == \"NY\":\n",
    "        return 20  # Return 20°C for New York\n",
    "\n",
    "    # Check if the city is Beijing ('BJ')\n",
    "    elif city == \"BJ\":\n",
    "        return 30  # Return 30°C for Beijing\n",
    "\n",
    "    # If the city is neither 'NY' nor 'BJ', return -1 to indicate unknown city\n",
    "    else:\n",
    "        return -1\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "weather_tool = FunctionTool.from_defaults(fn=get_weather)\n",
    "\n",
    "agent = ReActAgent.from_tools([multiply_tool, add_tool, weather_tool], llm=llm, verbose=True)\n",
    "\n",
    "response = agent.chat(\"纽约天气怎么样?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到模型的推理能力很强，将纽约转成了 NY。可以在 arize_phoenix 中看到 agent 的具体提示词，工具被装换成了提示词。\n",
    "ReActAgent 使得业务自动向代码转换成为可能，只要有 API 模型就可以调用，很多业务场景都适用，LlamaIndex 提供了一些开源的工具实现，可以到官网查看。\n",
    "\n",
    "虽然 Agent 可以实现业务功能， 但是一个 Agent 不能完成所有的功能，这也符合软件解耦的设计原则，不同的 Agent 可以完成不同的任务，各司其职，Agent 之间可以进行交互、通信，类似于微服务。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dbgpt_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
